import numpy as np
import scipy.sparse as sp
import torch
import pandas as pd
import matplotlib.pyplot as plt  # 新增：用于画图
import os
import h5py

def load_data(path, dataset=None,month=-1):
    """Load dataset"""
    # print('Loading {} dataset...'.format(dataset))
    
    # 生成文件名前缀
    base_path = fr"{path}stock/{dataset}/stock_{month}"
    
    # 读取.content文件
    idx_features_labels = np.genfromtxt(f"{base_path}.content", dtype=np.dtype(str))
    features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)
    labels = idx_features_labels[:, -1].astype(np.float32)

    # 构建图结构
    idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
    idx_map = {j: i for i, j in enumerate(idx)}
    edges_unordered = np.genfromtxt(f"{base_path}.cites", dtype=np.int32)

    # 修复关键：过滤掉不在 feature 节点列表中的边
    valid_nodes = set(idx)
    edges_filtered = [e for e in edges_unordered if e[0] in valid_nodes and e[1] in valid_nodes]
    edges_filtered = np.array(edges_filtered)

    # 映射索引（重新用 idx_map 转换成索引位置）
    edges = np.array(list(map(idx_map.get, edges_filtered.flatten())), dtype=np.int32).reshape(-1, 2)


    # edges = np.array(list(map(idx_map.get, edges_unordered.flatten())), dtype=np.int32).reshape(edges_unordered.shape)
    
    # 创建邻接矩阵
    adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])),
                        shape=(len(idx), len(idx)),
                        dtype=np.float32)

    # 对称化邻接矩阵
    adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)

    # 归一化处理
    features = normalize_features(features)
    # adj = normalize_adj(adj + sp.eye(adj.shape[0]))
    adj = adj + sp.eye(adj.shape[0])

    # 转换为PyTorch张量
    adj = torch.FloatTensor(np.array(adj.todense()))
    features = torch.FloatTensor(np.array(features.todense()))
    labels = torch.FloatTensor(labels).flatten()

    return adj, features, labels


def split_data_to_month(grouped_node, grouped_return, ret_col_name):
    date_arrays = {}
    for date, node_group in grouped_node:
        return_group = grouped_return.get_group(date)
        stock_indices = node_group.index.get_level_values('stock').values.reshape(-1, 1)
        combined_array = np.hstack([
            stock_indices,
            node_group.values,
            return_group[ret_col_name].values.reshape(-1, 1)
        ])
        date_arrays[date] = combined_array

    print('节点-标签拼接输出示例：')
    sample_key=list(date_arrays.keys())[0]
    display('第一个月数据：',date_arrays[sample_key])
    return date_arrays


# def save_month_data(date_arrays:dict, return_col, edge):
#     output_dir = f"stock/stock_node_{return_col}_edge_{edge}"
#     os.makedirs(output_dir, exist_ok=True)
#     for date, array in date_arrays.items():
#         file_date = str(date).replace("-", "")
#         filename = f"stock_{file_date}.content"
#         filepath = os.path.join(output_dir, filename)
        
#         # 生成格式字符串：第一列为 %d，其余列为 %.6f
#         n_columns = array.shape[1]
#         fmt_list = ["%d"] + ["%.6f"] * (n_columns - 1)  # 第一列整型，其他浮点型
#         fmt_str = "\t".join(fmt_list)  # 用制表符分隔
        
#         np.savetxt(
#             filepath,
#             array,
#             fmt=fmt_str,  # 自定义格式
#             delimiter="\t"
#         )
#     print('所有节点数据保存完毕')


def save_month_data(date_arrays:dict, adj_df: pd.DataFrame, 
                    ret_col_name, edge, is_rank = False, save_weight=False):
    """
    将邻接矩阵DataFrame按月保存为Cora的.cites格式文件
    
    参数:
        adj_df: 邻接矩阵DataFrame (MultiIndex: date, stock; Columns: stock)
        output_dir: 输出目录
        edge: 数据集名称(如'close'或'sec')
        save_weight: 是否保存权重
    """
    # 创建输出目录

    output_dir = f"stock/stock_node_{ret_col_name}_edge_{edge}"
    os.makedirs(output_dir, exist_ok=True)

    # 保存节点数据
    for date, array in date_arrays.items():
        file_date = str(date).replace("-", "")
        filename = f"stock_{file_date}.content"
        filepath = os.path.join(output_dir, filename)
        
        n_columns = array.shape[1]
        fmt_list = ["%d"] + ["%.6f"] * (n_columns - 1)  
        fmt_str = "\t".join(fmt_list)  
        
        np.savetxt(
            filepath,
            array,
            fmt=fmt_str,  
            delimiter="\t"
        )
    print('所有节点数据保存完毕')

    # 保存边数据
    total_months = len(adj_df.groupby(level='date'))
    files_saved = 0
    
    for date, date_group in adj_df.groupby(level='date'):
        date = pd.to_datetime(date)
        
        filename = f"stock_{date.strftime('%Y%m')}.cites"
        filepath = os.path.join(output_dir, filename)
        
        row_indices = date_group.index.get_level_values('stock')
        col_indices = date_group.columns
        
        matrix = date_group.values
        
        with open(filepath, 'w') as f:
            rows, cols = np.where(matrix == 1)
            
            for row_idx, col_idx in zip(rows, cols):
                stock_idx = row_indices[row_idx]
                connected_idx = col_indices[col_idx]
                
                if save_weight:
                    f.write(f"{stock_idx}\t{connected_idx}\t{1.000000}\n")
                else:
                    f.write(f"{stock_idx}\t{connected_idx}\n")
        
        files_saved += 1
    
    print(f'所有{edge}边数据保存完毕')


def describe_node_features(features):
    stats = {
        "节点数量": features.shape[0],
        "特征维度": features.shape[1],
        "各特征均值": np.mean(features, axis=0),
        "各特征标准差": np.std(features, axis=0),
        "各特征最小值": np.min(features, axis=0),
        "各特征最大值": np.max(features, axis=0),
        "特征稀疏度": (features == 0).mean()  
    }
    return pd.DataFrame(stats)


def zscore(x):
    return (x - x.mean()) / x.std(ddof=0)

# def normalize_adj(mx):
#     """Row-normalize sparse matrix"""
#     rowsum = np.array(mx.sum(1))
#     r_inv_sqrt = np.power(rowsum, -0.5).flatten()
#     r_inv_sqrt[np.isinf(r_inv_sqrt)] = 0.
#     r_mat_inv_sqrt = sp.diags(r_inv_sqrt)
#     return mx.dot(r_mat_inv_sqrt).transpose().dot(r_mat_inv_sqrt)


# def normalize_features(mx):
#     """Row-normalize sparse matrix"""
#     rowsum = np.array(mx.sum(1))
#     r_inv = np.power(rowsum, -1).flatten()
#     r_inv[np.isinf(r_inv)] = 0.
#     r_mat_inv = sp.diags(r_inv)
#     mx = r_mat_inv.dot(mx)
#     return mx